package com.test.demo.aspect;

import com.baomidou.mybatisplus.core.conditions.query.QueryWrapper;
import com.baomidou.mybatisplus.core.toolkit.StringUtils;
import com.test.demo.annotation.DataSource;
import com.test.demo.config.DataSourceType;
import com.test.demo.config.DynamicDataSourceContextHolder;
import com.test.demo.entity.e.DataSourceProperties;
import com.test.demo.service.DataSourceService;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Pointcut;
import org.aspectj.lang.reflect.MethodSignature;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.context.annotation.Lazy;
import org.springframework.core.annotation.AnnotationUtils;
import org.springframework.core.annotation.Order;
import org.springframework.stereotype.Component;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import javax.annotation.Resource;
import javax.servlet.http.HttpServletRequest;
import java.util.Objects;

/**
 * 多数据源处理
 */
@Aspect
@Order(1)
@Component
@Lazy
@Slf4j
public class DataSourceAspect {

    @Resource
    private DataSourceService dataSourceService;

    @Pointcut("@annotation(com.test.demo.annotation.DataSource)"
            + "|| @within(com.test.demo.annotation.DataSource)")
    public void dsPointCut() {
    }

    @Around("dsPointCut()")
    public Object around(ProceedingJoinPoint point) throws Throwable {
        DataSource dataSource = getDataSource(point);
        String dataSourceType;

        if (dataSource != null) {
            if (dataSource.value() == DataSourceType.NONE) {
                // 注解未指定数据源，从请求参数获取
                dataSourceType = determineDataSourceFromRequest();
            } else {
                // 使用注解指定的数据源
                dataSourceType = dataSource.value().name();
            }
        } else {
            // 没有注解，使用默认主库
            dataSourceType = DataSourceType.MASTER.name();
        }

        DynamicDataSourceContextHolder.setDataSourceKey(dataSourceType);
        log.info("Set datasource to: {}", dataSourceType);

        try {
            return point.proceed();
        } finally {
            DynamicDataSourceContextHolder.clearDataSourceKey();
        }
    }

    /**
     * 获取需要切换的数据源
     */
    public DataSource getDataSource(ProceedingJoinPoint point) {
        MethodSignature signature = (MethodSignature) point.getSignature();
        DataSource dataSource = AnnotationUtils.findAnnotation(signature.getMethod(), DataSource.class);
        if (Objects.nonNull(dataSource)) {
            return dataSource;
        }

        return AnnotationUtils.findAnnotation(signature.getDeclaringType(), DataSource.class);
    }

    /**
     * 从请求参数中获取数据源
     */
    private String determineDataSourceFromRequest() {
        ServletRequestAttributes attributes = (ServletRequestAttributes) RequestContextHolder.getRequestAttributes();
        if (attributes != null) {
            HttpServletRequest request = attributes.getRequest();
            String host = request.getHeader("Host");
            log.info("host:{}", host);
            DataSourceProperties hostObj = dataSourceService.getOne(new QueryWrapper<DataSourceProperties>().eq("host", host));
            if (null == hostObj) {
                throw new RuntimeException("host is invald.");
            }
            String dataSourceParam = hostObj.getDatasourceKey();
            if (StringUtils.isNotBlank(host) && StringUtils.isNotBlank(dataSourceParam)) {
                try {
                    return dataSourceParam;
                } catch (IllegalArgumentException e) {
                    log.warn("Invalid dataSource parameter: {}, using default MASTER", dataSourceParam);
                }
            } else {
                throw new RuntimeException("host is NULL");
            }
        }

        return DataSourceType.MASTER.name();
    }
} 